import pandas as pd
from statsmodels.tsa.statespace.sarimax import SARIMAX

# 假设CSV文件包含'发货时间'（日期时间格式）、'货品类型'和'销售量'列
df = pd.read_csv('clear_data.csv', parse_dates=['create_time'])
df.set_index('create_time', inplace=True)

# 创建月份列
df['月份'] = df.index.month

# 按货品类型和月份聚合销售量
grouped = df.groupby(['hplx', '月份'])['jz'].sum().reset_index()

# 准备进行预测的数据结构
predictions = {}

# 对每种产品进行预测
for product_type, product_data in grouped[grouped['hplx'].notnull()].groupby('hplx'):
    # 确保数据按月份排序
    product_data = product_data.sort_values('月份')

    # 提取销售量和月份
    y = product_data['jz']
    X = product_data['月份'] - 1  # 月份通常从1开始，但索引从0开始

    # 拟合模型（这里只是一个示例，参数可能需要调整）
    model = SARIMAX(y, order=(1, 1, 1), seasonal_order=(1, 1, 1, 12))
    results = model.fit()

    # 预测未来几个月的销售量
    # 注意：这里假设我们预测未来12个月的销售量
    start = len(y)
    end = start + 11
    forecast = results.get_forecast(steps=12)
    forecast_sales = forecast.predicted_mean

    # 生成预测的月份（从最后一个已知月份开始）
    last_month = X.iloc[-1]
    forecast_months = [last_month + i + 1 for i in range(12)]  # 加1是因为月份从1开始

    # 将预测结果保存为Series，索引为月份
    predictions[product_type] = pd.Series(forecast_sales, index=forecast_months)

# 将预测结果转换为DataFrame
prediction_df = pd.DataFrame(predictions).T

# 将DataFrame保存到CSV文件中
prediction_df.to_csv('predicted_sales_with_seasonality.csv')

# 打印消息以确认保存
print("预测结果已保存到predicted_sales_with_seasonality.csv文件中。")